import argparse
import re
from pathlib import Path

import PyPDF2


def is_likely_end_of_paragraph(line: str) -> bool:
    """
    Heuristically checks if a line is likely the end of a paragraph.
    """
    if not line:
        return True
    # Ends with a period, question mark, exclamation mark, or colon.
    if line.strip().endswith((".", "?", "!", ":")):
        return True
    # Is a short line (e.g., a heading or a small sentence)
    if len(line.strip()) < 60:
        return True
    return False


def convert_pdf_to_markdown(pdf_path: Path, output_dir: Path):
    """
    Converts a single PDF file to a markdown file, attempting to merge paragraphs
    using more robust heuristics.

    Args:
        pdf_path (Path): The path to the PDF file.
        output_dir (Path): The directory to save the markdown file.
    """
    try:
        print(f"Processing {pdf_path}...")
        with open(pdf_path, "rb") as pdf_file:
            pdf_reader = PyPDF2.PdfReader(pdf_file)
            final_text_parts = []

            for page in pdf_reader.pages:
                page_text = page.extract_text()
                if not page_text or not page_text.strip():
                    continue

                # Clean up the text first
                # Handle words hyphenated at the end of a line
                page_text = re.sub(r"-\s*\n", "", page_text)

                lines = page_text.split("\n")
                paragraphs = []
                current_paragraph = ""

                for i, line in enumerate(lines):
                    clean_line = line.strip()
                    if not clean_line:
                        if current_paragraph:
                            paragraphs.append(current_paragraph)
                            current_paragraph = ""
                        continue

                    if not current_paragraph:
                        current_paragraph = clean_line
                    # Check if the line should be appended to the current paragraph
                    # or start a new one
                    elif not is_likely_end_of_paragraph(current_paragraph.split("\n")[-1]):
                        current_paragraph += " " + clean_line
                    else:
                        paragraphs.append(current_paragraph)
                        current_paragraph = clean_line

                if current_paragraph:
                    paragraphs.append(current_paragraph)

                if paragraphs:
                    final_text_parts.append("\n\n".join(paragraphs))

            if not final_text_parts:
                print(
                    f"Warning: No text extracted from {pdf_path}. It might be an image-based PDF or have other issues."
                )
                return

            final_text = "\n\n---\n\n".join(final_text_parts)

            # Create output directory if it doesn't exist
            output_dir.mkdir(parents=True, exist_ok=True)

            # Save the text to a markdown file
            output_filename = pdf_path.stem + ".md"
            output_filepath = output_dir / output_filename

            with open(output_filepath, "w", encoding="utf-8") as md_file:
                md_file.write(final_text)

            print(f"Successfully converted {pdf_path} to {output_filepath}")

    except PyPDF2.errors.PdfReadError:
        print(
            f"Error: Could not read {pdf_path}. It may be corrupted or not a valid PDF."
        )
    except Exception as e:
        print(f"An unexpected error occurred while processing {pdf_path}: {e}")


def main():
    """
    Main function to handle command-line arguments and start the conversion process.
    """
    parser = argparse.ArgumentParser(
        description="Convert PDF files to Markdown. Extracts text from PDFs and formats it into paragraphs."
    )
    parser.add_argument(
        "input_path",
        type=Path,
        help="Path to a PDF file or a directory containing PDF files.",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        default="output_markdown",
        help="Path to the output directory for markdown files. Defaults to 'output_markdown'.",
    )

    args = parser.parse_args()
    input_path: Path = args.input_path
    output_dir: Path = args.output

    if not input_path.exists():
        print(f"Error: The path '{input_path}' does not exist.")
        return

    if input_path.is_file():
        if input_path.suffix.lower() == ".pdf":
            convert_pdf_to_markdown(input_path, output_dir)
        else:
            print(f"Error: The file '{input_path}' is not a PDF file.")
    elif input_path.is_dir():
        pdf_files = list(input_path.glob("*.pdf"))
        if not pdf_files:
            print(f"No PDF files found in directory '{input_path}'.")
            return

        for pdf_path in pdf_files:
            convert_pdf_to_markdown(pdf_path, output_dir)
    else:
        print(f"Error: The path '{input_path}' is not a valid file or directory.")


if __name__ == "__main__":
    main() 